#!/usr/bin/python
'''
Desc: 基于特定任务生成大模型prompt,参考https://zhuanlan.zhihu.com/p/645045642
'''
from enum import Enum

class TaskType(Enum):
    GENERATE = 1
    SUMMARIZE = 2
    CLASSIFY = 3
    EXTRACT = 4
    SEARCH = 5
    REWRITE = 6

class Task:
    @classmethod
    def generate(cls, data):
        return cls.TEMPLATE.format(text=data)

class GenerateTask(Task):
    TEMPLATE = '''
请对以下文本进行总结,以一个5岁孩子能听懂的方式进行回答.
{text}
'''

class SummarizeTask(Task):
    TEMPLATE = '''
请对以下文本进行总结,以可爱的方式进行回答.
{text}
'''

class ClassifyTask(Task):
    TEMPLATE = '''
请完成情感分类任务,给定一个句子,从['negative','positive']中分配一个标签,只返回标签不要返回其他任何文本.

Sentence: 这真是太有趣了.
Label:positive
Sentence: 这件衣服的材质有点差.
Label:negative

{text}
Label:
'''

class ExtractTask(Task):
    TEMPLATE = '''
请完成关键词提取任务,给定一个句子,从中提取水果名称,如果文中没有水果请回答“文中没有提到水果”.不要回答其他无关内容.

Sentence: 在果摊上,摆放着各式水果.成熟的苹果,香甜的香蕉,翠绿的葡萄,以及紫色的蓝莓.
fruit names: 苹果,香蕉,葡萄,蓝莓

{text}
fruit names:
'''

class SearchTask(Task):
    # use embedding not support prompt engineering
    TEMPLATE = ""

class RewriteTask(Task):
    TEMPLATE = '''
请完成文本纠错的任务,给定一段文本,对文本中的错别字或语法错误进行修改,并返回正确的版本,如果文本中没有错误,什么也不要返回.

text: 黄昏，一缕轻烟从烟囱里请缨地飘出来，地面还特么的留有一丝余热，如果说正午像精力允沛的青年，那黄昏就像幽雅的少女，清爽的风中略贷一丝暖意。
correct: 黄昏，一缕轻烟从烟囱里轻轻地飘出来，地面还留有一丝余热，如果说正午像精力充沛的青年，那黄昏就像优雅的少女，清爽的风中略带一丝暖意。
text: 胎头望着天空，只见红彤彤的晚霞己经染红大半片天空了，形状更是千资百态。
correct: 抬头望着天空，只见红彤彤的晚霞己经染红大半片天空了，形状更是千姿百态。

{text}
correct:
'''

def get_task_text(mode, text):
    task_generator = {
        TaskType.GENERATE : GenerateTask,
        TaskType.CLASSIFY : ClassifyTask,
        TaskType.SUMMARIZE: SummarizeTask,
        TaskType.SEARCH: SearchTask,
        TaskType.EXTRACT: ExtractTask,
        TaskType.REWRITE: RewriteTask
    }
    assert mode in task_generator, f"{mode} should in task_generator"
    output = task_generator[mode].generate(text)
    return output

if __name__ == '__main__':
    #output = get_task_text(TaskType.CLASSIFY, "他刚才说了一堆废话.")
    text = """
草莓、蓝莓、香蕉和橙子等水果富含丰富的营养素，包括维生素、纤维和抗氧化剂，对于维持健康和预防疾病有重要作用。
"""
    output = get_task_text(TaskType.EXTRACT, text)

    print(output)

